//! Sled database operations.

use std::time;

use rand::RngCore;
use sled::transaction::{ConflictableTransactionError, TransactionalTree};
use sled::Tree;

use aspect_db_mq_types::{Error, QueueId, ReaderId, ID_LEN};
use aspect_msgqueue::{Msg, QueueView};

use super::errors::*;
use super::state::QueueState;

pub struct LlmqDatabase {
    db: Tree,
}

impl LlmqDatabase {
    pub fn new(db: Tree) -> Result<LlmqDatabase, Error> {
        Ok(LlmqDatabase { db })
    }

    /// Generates a new random queue ID.  Does not exist yet.
    fn gen_queue_id(&self) -> QueueId {
        let mut buf = [0u8; ID_LEN];
        let mut rng = rand::rngs::OsRng;
        rng.fill_bytes(&mut buf);
        QueueId::from(buf)
    }

    /// Generates a new random reader ID.  Does not exist anywhere yet.
    pub fn gen_reader_id(&self) -> ReaderId {
        ReaderId::new_rand()
    }

    #[cfg(test)]
    fn get_queue_state(&self, queue_id: &QueueId) -> Result<Option<QueueState>, DbError> {
        let qk = get_queue_state_key(queue_id);

        let Some(raw) = self.db.get(&qk)? else {
            return Ok(None);
        };

        let state = aspect_codec::decode::<QueueState>(&raw)?;
        state.sanity_check()?;
        Ok(Some(state))
    }

    /// Creates a new MQ with a set of initial reader IDs.
    pub fn create_queue(&self, readers: Vec<ReaderId>) -> Result<QueueId, DbError> {
        // Create the new queue and insert it.
        let qid = self.gen_queue_id();
        let qk = get_queue_state_key(&qid);
        let state = QueueState::new_with_readers(readers);
        state.sanity_check().expect("mq: inited state invalid");
        let raw = aspect_codec::encode_to_vec(&state)?;

        self.db.insert(&qk, raw)?;

        Ok(qid)
    }

    /// Subscribes a reader ID to a queue ID, returning `true` if added or
    /// `false` if the reader was already subscribed.
    pub fn subscribe_reader_to_queue(
        &self,
        queue_id: &QueueId,
        reader_id: &ReaderId,
    ) -> Result<bool, DbError> {
        let succ = self.db.transaction(move |tt| {
            let mut state = load_queue_state_txly(tt, queue_id)?;
            let succ = state.try_insert_reader(reader_id);

            // If we changed something, then write back the change.
            if succ {
                save_queue_state_txly(tt, queue_id, &state)?;
            }

            Ok(succ)
        })?;

        Ok(succ)
    }

    /// Subscribes a new reader ID to a queue, like `subscribe_reader_to_queue`,
    /// but generates a fresh reader ID so should always succesed.
    pub fn subscribe_new_reader_to_queue(&self, queue_id: &QueueId) -> Result<ReaderId, DbError> {
        let rid = self.gen_reader_id();
        let succ = self.subscribe_reader_to_queue(queue_id, &rid)?;
        assert!(
            succ,
            "mq: unsuccessful in subscribing fresh reader ID to queue"
        );
        Ok(rid)
    }

    /// Unsubscribes a reader ID from a queue ID, returning `true` if removed or
    /// `false` if the reader was not subscribed.
    pub fn unsubscribe_reader_from_queue(
        &self,
        queue_id: &QueueId,
        reader_id: &ReaderId,
    ) -> Result<bool, DbError> {
        let succ = self.db.transaction(move |tt| {
            let mut state = load_queue_state_txly(tt, queue_id)?;
            let succ = state.try_remove_reader(reader_id);

            // If we changed something, then write back the change.
            if succ {
                save_queue_state_txly(tt, queue_id, &state)?;
            }

            Ok(succ)
        })?;

        Ok(succ)
    }

    /// Inserts a new message into the queue database, returning the index if
    /// successful.
    pub fn insert_msg_into_queue(&self, queue_id: &QueueId, msg: &Msg) -> Result<u64, DbError> {
        let res = self.db.transaction(move |tt| {
            // Get the current queue state.
            let mut state = load_queue_state_txly(tt, queue_id)?;

            // Figure out the next ID.
            let idx = state.advance_idx();
            let mk = get_msg_key(queue_id, idx);

            // Insert the message.
            let msgbuf = aspect_codec::encode_to_vec(msg).map_err(into_cte)?;
            tt.insert(&mk, msgbuf)?;

            // Insert the new queue state.
            save_queue_state_txly(tt, queue_id, &state)?;

            Ok(idx)
        })?;

        Ok(res)
    }

    /// Takes a generic message payload and inserts it into the queue at the current time.
    pub fn submit_generic_msg(&self, queue_id: &QueueId, body: Vec<u8>) -> Result<u64, DbError> {
        let now = time::Instant::now().elapsed().as_millis() as u64;
        let m = Msg::new(now, body);
        self.insert_msg_into_queue(queue_id, &m)
    }

    /// Queries the view coordinates for a queue, maybe from the perspective of
    /// a particular reader.
    pub fn query_queue_view(
        &self,
        queue_id: &QueueId,
        reader_id: Option<&ReaderId>,
    ) -> Result<QueueView, DbError> {
        let qk = get_queue_state_key(queue_id);

        let Some(raw) = self.db.get(&qk)? else {
            return Err(Error::UnknownQueue(*queue_id).into());
        };

        let state = aspect_codec::decode::<QueueState>(&raw)?;
        state.sanity_check()?;

        let rp = match reader_id {
            Some(rid) => state
                .get_reader_position(rid)
                .map_err(|_| Error::UnknownQueue(*queue_id))?,
            None => u64::MAX,
        };

        Ok(QueueView::new(state.next_idx(), state.tail_idx(), rp))
    }

    /// Queries a range of messages from the queue, if there are any in the range.
    /// This is interpreted as a half-open interval.
    pub fn query_queue_message_range(
        &self,
        queue_id: &QueueId,
        reader_id: &ReaderId,
        min: u64,
        max: u64,
    ) -> Result<Vec<(u64, Msg)>, DbError> {
        // Checks to avoid starting a transaction if we know we won't get anything anyways.
        if min > max {
            return Err(Error::InvalidRange(min, max).into());
        }

        if min == max {
            return Ok(Vec::new());
        }

        let res = self.db.transaction(move |tt| {
            let state = load_queue_state_txly(tt, queue_id)?;

            // Make sure the reader actually has permission to read, and then
            // also ensure they haven't already declared that they've consumed a
            // message they're querying for.
            let Ok(consumed_pos) = state.get_reader_position(reader_id) else {
                return Err(into_cte(Error::UnknownQueue(*queue_id)));
            };

            let real_min = u64::max(min, consumed_pos);

            let mut msgs = Vec::new();

            // More checks to avoid making unnecesary queries.
            if max <= state.tail_idx() {
                return Ok(msgs);
            }

            if real_min >= state.next_idx() {
                return Ok(msgs);
            }

            // Also add some more bounding.
            let real_max = u64::min(max, state.next_idx());

            for i in real_min..real_max {
                let mk = get_msg_key(queue_id, i);

                let Some(msg_raw) = tt.get(&mk)? else {
                    continue;
                };

                // TODO should we move the parsing outside the loop?
                let msg = aspect_codec::decode(&msg_raw).map_err(into_cte)?;
                msgs.push((i, msg));
            }

            Ok(msgs)
        })?;

        Ok(res)
    }

    /// Updates the position a reader declares as being fully consumed, updating
    /// the queue state as appropriate.  The target specified here will now be
    /// the index of the earliest message that they can still read, wether or
    /// not it has been written to the queue or not.
    pub fn update_consumed_position(
        &self,
        queue_id: &QueueId,
        reader_id: &ReaderId,
        target: u64,
    ) -> Result<(u64, u64), DbError> {
        let res = self.db.transaction(move |tt| {
            let mut state = load_queue_state_txly(tt, queue_id)?;

            let r = state
                .update_reader_position(reader_id, target)
                .map_err(|_| Error::UnknownQueue(*queue_id))
                .map_err(into_cte)?;

            // TODO purge old messages somewhere, maybe this should happen
            // lazily so it doesn't block other transactions if there's many
            // entries to delete?
            state.refresh_tail();

            save_queue_state_txly(tt, queue_id, &state)?;

            Ok(r)
        })?;

        Ok(res)
    }
}

fn get_queue_state_key(id: &QueueId) -> [u8; 17] {
    let mut k = [0u8; 17];
    k[0] = b'Q';
    k[1..].copy_from_slice(id.as_ref());
    k
}

fn get_msg_key(id: &QueueId, idx: u64) -> [u8; 25] {
    let mut k = [0u8; 25];
    k[0] = b'M';
    k[1..17].copy_from_slice(id.as_ref());
    k[17..].copy_from_slice(&idx.to_be_bytes()[..]);
    k
}

/// Loads a queue's state structure and validates it, such that it can be done
/// in a database transaction.
fn load_queue_state_txly(
    tt: &TransactionalTree,
    queue_id: &QueueId,
) -> Result<QueueState, ConflictableTransactionError<Error>> {
    let qk = get_queue_state_key(queue_id);

    let Some(state_raw) = tt.get(&qk)? else {
        return Err(into_cte(Error::UnknownQueue(*queue_id)));
    };

    let state = aspect_codec::decode::<QueueState>(&state_raw).map_err(into_cte)?;
    state.sanity_check().map_err(into_cte)?;

    Ok(state)
}

/// Saves a queue's state, validating it, like the reverse of `load_queue_state_txly`.
fn save_queue_state_txly(
    tt: &TransactionalTree,
    queue_id: &QueueId,
    state: &QueueState,
) -> Result<(), ConflictableTransactionError<Error>> {
    // TODO somehow make this optional?
    state.sanity_check().map_err(into_cte)?;

    let state_raw = aspect_codec::encode_to_vec(state).map_err(into_cte)?;
    let qk = get_queue_state_key(queue_id);
    tt.insert(&qk, state_raw)?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::LlmqDatabase;

    #[test]
    fn test_mq_write_read_consume() {
        let td = tempfile::TempDir::new().expect("test: init tempdir");
        let db = sled::open(td.path()).expect("test: open db");
        let tree = db.open_tree("queuedata").expect("test: open tree");
        let mq = LlmqDatabase::new(tree).expect("test: create mqds");

        let r1 = mq.gen_reader_id();
        let r2 = mq.gen_reader_id();
        let r3 = mq.gen_reader_id();

        // Create the queue.
        let qid = mq
            .create_queue(vec![r1, r2, r3])
            .expect("test: create queue");

        // Insert some messages.
        let m1 = mq
            .submit_generic_msg(&qid, vec![1, 2, 3])
            .expect("test: submit");
        let m2 = mq
            .submit_generic_msg(&qid, vec![3, 4, 5])
            .expect("test: submit");
        let m3 = mq
            .submit_generic_msg(&qid, vec![5, 6, 7])
            .expect("test: submit");

        // Dump some stuff.
        let qs1 = mq.get_queue_state(&qid).expect("test: get state");
        let qv1 = mq
            .query_queue_view(&qid, Some(&r1))
            .expect("test: get view");
        eprintln!("m1 {m1}\nqv1 {qv1:?}\nqs1 {qs1:#?}");

        // Query the messages back.
        let r1q1 = mq
            .query_queue_message_range(&qid, &r1, 0, 100)
            .expect("test: get range");
        eprintln!("r1q1 {r1q1:#?}");

        // Now consume 2 messages and see that the states make sense.
        mq.update_consumed_position(&qid, &r1, 1)
            .expect("test: consume");

        // Check that the new limits make sense.
        let qs2 = mq.get_queue_state(&qid).expect("test: get state");
        let r1q2 = mq
            .query_queue_message_range(&qid, &r1, 0, 100)
            .expect("test: get range");
        eprintln!("qs2 {qs2:#?}");
        eprintln!("r1q2 {r1q2:#?}");
    }
}
